import logging

import MySQLdb
from MySQLdb.cursors import CursorUseResultMixIn, CursorTupleRowsMixIn, BaseCursor
import psycopg2
import pyarrow as pa

from src.py_data_sync.format import MYSQL_TO_ARROW_TYPES, PSYCOPG_TO_ARROW_TYPES


class IterCursor(CursorUseResultMixIn, CursorTupleRowsMixIn, BaseCursor):
    """
    use CursorUseResultMixIn to iterate result set rows
    """


class AbsSource:

    def execute(self):
        """
        execute sql query, generate data set schema for parquet file
        """
        raise NotImplementedError

    def get_sql(self):
        raise NotImplementedError

    def __iter__(self):
        return self


class MysqlSource(AbsSource):

    def __init__(self, conf) -> None:
        self.conf = conf

    def execute(self):
        self.sql = self.get_sql()
        self._conn = self.__create_db_conn()
        self._cursor = self._conn.cursor()
        self._cursor.execute(self.sql)
        self.desc = self._cursor.description
        self.schema = self.__get_schema()

    def get_sql(self):
        columns = "*"
        if self.conf["detail"].get("columns"):
            columns = "{},{}".format(
                self.conf["detail"]["primary_key"], self.conf["detail"]["columns"]
            )
        limit = self.conf["detail"].get("limit")
        table = self.conf["detail"]["table"]

        sql = f"select {columns} from {table}"
        if limit and limit > 0:
            sql += f" limit {str(limit)}"

        logging.info("execute sql: %s" % sql)
        return sql

    def __create_db_conn(self):
        self.conf["connect"]["cursorclass"] = IterCursor
        return MySQLdb.connect(**self.conf["connect"])

    def __get_schema(self):
        fields = []
        desc = self._cursor.description
        for d in desc:
            field_name = d[0]
            mysql_type = d[1]
            arrow_type = MYSQL_TO_ARROW_TYPES[mysql_type]
            fields.append(pa.field(field_name, arrow_type))
        return pa.schema(fields)

    def __next__(self):
        rows = self._cursor.fetchmany(1000)
        if not rows:
            self._cursor.close()
            self._conn.close()
            raise StopIteration
        return rows


class PgsqlSource(AbsSource):

    def __init__(self, conf) -> None:
        self.conf = conf

    def execute(self):
        self.sql = self.get_sql()
        self._conn = self.__create_db_conn()
        self._cursor = self._conn.cursor()
        self._cursor.execute(self.sql)
        self.desc = self._cursor.description
        self.schema = self.__get_schema()

    def get_sql(self):
        columns = "*"
        if self.conf["detail"].get("columns"):
            columns = "{},{}".format(
                self.conf["detail"]["primary_key"], self.conf["detail"]["columns"]
            )
        limit = self.conf["detail"].get("limit")
        table = self.conf["detail"]["table"]

        sql = f"select {columns} from {table}"
        if limit and limit > 0:
            sql += f" limit {str(limit)}"

        logging.info("execute sql: %s" % sql)

        return sql

    def __create_db_conn(self):
        return psycopg2.connect(**self.conf["connect"])

    def __get_schema(self):
        fields = []
        desc = self._cursor.description
        for d in desc:
            field_name = d.name
            pgsql_type = d.type_code
            arrow_type = PSYCOPG_TO_ARROW_TYPES[pgsql_type]
            fields.append(pa.field(field_name, arrow_type))
        return pa.schema(fields)

    def __next__(self):
        rows = self._cursor.fetchmany(1000)
        if not rows:
            self._cursor.close()
            self._conn.close()
            raise StopIteration
        return rows
